import numpy as np
from sklearn.metrics import precision_score, recall_score, f1_score

# 1. 模拟检索结果和真实标签
def generate_mock_data():
    """ 
    模拟检索结果和真实标签
    :return: 检索结果和真实标签
    """ 
    np.random.seed(42)
    relevant_items = np.random.choice([0, 1], size=100, p=[0.8, 0.2])  # 模拟100个向量，其中20%为相关项
    retrieved_items = np.random.choice([0, 1], size=100, p=[0.7, 0.3])  # 模拟系统返回结果
    return relevant_items, retrieved_items

relevant_items, retrieved_items = generate_mock_data()

# 2. 计算精度、召回率和F1评分
def calculate_metrics(true_labels, predicted_labels):
    """ 
    计算精度、召回率和F1评分
    :param true_labels: 真实标签
    :param predicted_labels: 检索结果标签
    :return: 精度、召回率、F1评分
    """ 
    precision = precision_score(true_labels, predicted_labels)
    recall = recall_score(true_labels, predicted_labels)
    f1 = f1_score(true_labels, predicted_labels)
    return precision, recall, f1

precision, recall, f1 = calculate_metrics(relevant_items, retrieved_items)

# 3. 输出混淆矩阵和结果
def print_confusion_matrix(true_labels, predicted_labels):
    """ 
    输出混淆矩阵
    :param true_labels: 真实标签
    :param predicted_labels: 检索结果标签
    """ 
    tp = np.sum((true_labels == 1) & (predicted_labels == 1))  # True Positive
    fp = np.sum((true_labels == 0) & (predicted_labels == 1))  # False Positive
    fn = np.sum((true_labels == 1) & (predicted_labels == 0))  # False Negative
    tn = np.sum((true_labels == 0) & (predicted_labels == 0))  # True Negative

    print("混淆矩阵:")
    print(f"TP (正确检索相关项): {tp}")
    print(f"FP (错误检索非相关项): {fp}")
    print(f"FN (未检索相关项): {fn}")
    print(f"TN (正确未检索非相关项): {tn}")

print_confusion_matrix(relevant_items, retrieved_items)

# 4. 输出精度、召回率和F1评分
print("\n性能指标:")
print(f"精度 (Precision): {precision:.4f}")
print(f"召回率 (Recall): {recall:.4f}")
print(f"F1评分 (F1 Score): {f1:.4f}")